from typing import List, Dict
from pyvis.network import Network
from datetime import datetime

class KnowledgeVisualizer:
    def __init__(self, rag_system=None):
        self.rag = rag_system
        
    def generate_graph(self, nodes: List[Dict]):
        """生成知识图谱可视化"""
        net = Network(height='750px', width='100%', bgcolor='#ffffff', font_color='#333')
        
        # 添加节点
        for node in nodes:
            net.add_node(
                node.get('id'),
                label=node.get('title', 'Unknown'),
                title=node.get('metadata', {}).get('title', ''),
                group=node.get('domain', 'default')
            )
            
        # 添加边
        for node in nodes:
            for ref in node.get('references', []):
                if any(n.get('id') == ref for n in nodes):
                    net.add_edge(node['id'], ref, title='cites')
        
        # 设置物理布局
        net.set_options("""
        const options = {
            "physics": {
                "forceAtlas2Based": {
                    "springLength": 100
                },
                "minVelocity": 0.75,
                "solver": "forceAtlas2Based"
            }
        }
        """)
        
        # 保存为HTML文件
        net.show('knowledge_graph.html')
        
    def create_timeline(self, papers: List[Dict]) -> str:
        """生成技术演进路线图"""
        # 按年份对论文进行分组
        papers_by_year = {}
        for paper in papers:
            year = paper.get('year', str(datetime.now().year))
            if year not in papers_by_year:
                papers_by_year[year] = []
            papers_by_year[year].append(paper)
        
        # 生成Mermaid图表
        mermaid = "```mermaid\ngantt\ntitle 技术演进路线\ndateFormat YYYY\n\n"
        
        # 添加时间线
        for year, year_papers in sorted(papers_by_year.items()):
            mermaid += f"section {year}\n"
            for paper in year_papers:
                title = paper.get('title', 'Unknown Paper')
                # 限制标题长度
                if len(title) > 30:
                    title = title[:27] + '...'
                mermaid += f"{title} : {year}\n"
        
        mermaid += "```"
        return mermaid
    
    def generate_trend_chart(self, papers: List[Dict]) -> str:
        """生成研究趋势图"""
        # 统计每年的论文数量
        papers_count = {}
        for paper in papers:
            year = paper.get('year', str(datetime.now().year))
            papers_count[year] = papers_count.get(year, 0) + 1
        
        # 生成趋势图（使用Mermaid）
        chart = "```mermaid\npie\ntitle 研究趋势分布\n"
        
        for year, count in sorted(papers_count.items()):
            chart += f"\"{year}\" : {count}\n"
        
        chart += "```"
        return chart